# encoding=utf-8
import re
import sys

from dalchemy.data import TextHelper
from dalchemy.llms import AzureLLM, OpenAILLM,ZhipuLLM
from tqdm import tqdm

helper = TextHelper()
PROMPT_Q = "根据上述文本中与医疗领域相关的内容与逻辑关系提出几个中文问题，注意，提出的问题应该提供充实的内容，使问题具有挑战性。"


def main():
    global llm
    infile = sys.argv[1]
    outfile = sys.argv[2]
    # infile = "data/passages.json"
    # outfile = "data/query_ls.json"
    passage_ls = helper.read_json(infile)
    print(len(passage_ls))
    res = []
    seen = set()
    qid = 0
    for passage in tqdm(passage_ls):
        # 构造prompt
        content = passage["content"]
        prompt = f"{content}\n\n{PROMPT_Q}"
        # 生成问题
        response = llm(prompt)
        print("response",response)
        # 抽取问题
        query_ls = response.replace("\\n", "\n").split("\n")
        pattern = r'^\d+\.\s*(.*)$'
        for query in query_ls:
            query = re.sub(pattern, r'\1', query)
            query = query.strip()
            if query and query not in seen:
                seen.add(query)
                record = {"qid": str(qid), "query": query, "pid": passage["pid"]}
                res.append(record)
                qid += 1

    helper.write_json(res, json_outfile=outfile)


if __name__ == '__main__':
    # azure_cfg = {
    #     "engine": "gpt-35-turbo-16k",
    #     "api_type": "azure",
    #     "api_base": "https://your_base",
    #     "api_version": "2023-05-15",
    #     "api_key": "your_key",
    #     "max_tokens": 8192
    # }
    # llm = AzureLLM(**azure_cfg)

    cfg = {
        "engine": "chatglm_pro",
        "api_key": "your_zhipu_key"
    }
    llm = ZhipuLLM(**cfg)

    main()
